import asyncio
import logging
from datetime import datetime, timedelta
from typing import Optional, List, Dict, Any, Union

from sqlalchemy import Column, String, Integer, DateTime, Text, Boolean, delete, BigInteger, JSON
from sqlalchemy import select
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker, async_scoped_session
from sqlalchemy.orm import declarative_base
from sqlalchemy.exc import SQLAlchemyError, DBAPIError, DisconnectionError

# 使用新的声明式基类
DeclarativeBase = declarative_base()

# 添加日志配置
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger('MessageDB')

class Message(DeclarativeBase):
    __tablename__ = 'messages'
    id = Column(Integer, primary_key=True, autoincrement=True)
    msg_id = Column(BigInteger, index=True, comment='消息唯一ID（整型）')
    new_msg_id = Column(BigInteger, index=True, comment='新消息ID（撤回等场景）', default=0)
    sender_wxid = Column(String(40), index=True, comment='消息发送人wxid')
    sender_nickname = Column(String(40), index=True, comment='消息发送人昵称')
    from_wxid = Column(String(40), index=True, comment='消息来源wxid')
    from_nickname = Column(String(40), index=True, comment='消息来源名称')
    msg_type = Column(Integer, comment='消息类型（整型编码）')
    content = Column(Text, comment='消息内容')
    org_content = Column(Text, comment='原始消息内容')
    timestamp = Column(DateTime, default=datetime.now, index=True, comment='消息时间戳')
    is_group = Column(Boolean, default=False, comment='是否群消息')
    at_names = Column(String(100), comment='at人信息')
    is_at = Column(Boolean, default=False, comment='是否at当前用户')

    def get(self, param, param1):
        return getattr(self, param, param1)


class GroupInfo(DeclarativeBase):
    __tablename__ = 'group_info'

    id = Column(Integer, primary_key=True, autoincrement=True)
    qid = Column(String(40), unique=True, index=True, comment='群ID')
    group_name = Column(String(100), comment='群名称')
    owner_wxid = Column(String(40), comment='群主wxid')
    member_count = Column(Integer, default=0, comment='群成员数量')
    announcement = Column(Text, comment='群公告')
    member_list = Column(JSON, comment='群成员列表')
    last_update = Column(DateTime, default=datetime.now, comment='最后更新时间')
    conversation = Column(JSON, comment='会话')


class FriendInfo(DeclarativeBase):
    __tablename__ = 'friend_info'

    id = Column(Integer, primary_key=True, autoincrement=True)
    wxid = Column(String(40), unique=True, index=True, comment='好友wxid')
    nickname = Column(String(100), comment='昵称')
    remark = Column(String(100), comment='备注')
    avatar_url = Column(String(500), comment='头像URL')
    phone = Column(String(20), comment='手机号')
    email = Column(String(100), comment='邮箱')
    is_friend = Column(Boolean, default=True, comment='是否为好友')
    last_update = Column(DateTime, default=datetime.now, comment='最后更新时间')
    conversation = Column(JSON, comment='会话')


class ChatRecord(DeclarativeBase):
    __tablename__ = 'chat_records'

    id = Column(Integer, primary_key=True, autoincrement=True)
    chat_id = Column(String(40), index=True, comment='聊天对象ID（个人或群）')
    chat_type = Column(String(10), comment='聊天类型：private/group')
    msg_id = Column(BigInteger, index=True, comment='消息ID')
    sender_wxid = Column(String(40), index=True, comment='发送者wxid')
    content = Column(Text, comment='消息内容')
    msg_type = Column(Integer, comment='消息类型')
    timestamp = Column(DateTime, default=datetime.now, index=True, comment='消息时间')
    is_read = Column(Boolean, default=False, comment='是否已读')

class ChatRecord(DeclarativeBase):
    __tablename__ = 'chatroom'

    chatroom_id = Column(String(20), primary_key=True, nullable=False, unique=True, index=True, autoincrement=False,
                             comment='chatroom_id')
    members = Column(JSON, nullable=False, default=list, comment='members')
    llm_thread_id = Column(JSON, nullable=False, default=lambda: {}, comment='llm_thread_id')



class MessageDB:
    _instance = None
    _lock = asyncio.Lock()

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
            cls._instance.initialized = False
            cls._instance.cleanup_task = None
            
            # 使用SQLite数据库
            db_url = "sqlite+aiosqlite:///xbot_messages.db"
            
            # 创建异步引擎
            cls._instance.engine = create_async_engine(
                db_url,
                echo=False,
                future=True,
                pool_pre_ping=True,  # 连接前进行ping检测
                pool_recycle=3600,   # 连接回收时间，1小时
                pool_timeout=30      # 连接超时时间
            )
            
            # 创建异步会话工厂
            async_session_factory = async_sessionmaker(
                cls._instance.engine,
                expire_on_commit=False
            )
            
            # 创建作用域会话
            cls._instance.async_session_factory = async_scoped_session(
                async_session_factory,
                scopefunc=asyncio.current_task
            )
            
        return cls._instance

    async def initialize(self):
        """异步初始化数据库"""
        if self.initialized:
            return
            
        async with self._lock:
            if self.initialized:
                return
                
            try:
                async with self.engine.begin() as conn:
                    await conn.run_sync(DeclarativeBase.metadata.create_all)
                
                # 启动定期清理任务
                self.cleanup_task = asyncio.create_task(self._periodic_cleanup())
                
                self.initialized = True
                logger.info("数据库初始化完成")
            except Exception as e:
                logger.error(f"数据库初始化失败: {e}")
                raise

    async def _periodic_cleanup(self):
        """定期清理过期消息"""
        while True:
            try:
                await asyncio.sleep(86400)  # 每24小时执行一次
                await self.cleanup_messages()
                logger.info("执行了定期消息清理")
            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(f"定期清理任务异常: {e}")
                await asyncio.sleep(3600)  # 出错后1小时后重试

    async def ensure_connection(self):
        """确保数据库连接有效，如果无效则重新初始化"""
        try:
            # 尝试执行一个简单查询来验证连接
            async with self.async_session_factory() as session:
                await session.execute(select(1))
            return True
        except (DBAPIError, DisconnectionError) as e:
            logger.error(f"数据库连接错误，尝试重新初始化: {e}")
            try:
                # 重新创建引擎和会话工厂
                db_url = "sqlite+aiosqlite:///xbot_messages.db"
                self.engine = create_async_engine(
                    db_url,
                    echo=False,
                    future=True,
                    pool_pre_ping=True,
                    pool_recycle=3600,
                    pool_timeout=30
                )
                
                async_session_factory = async_sessionmaker(
                    self.engine,
                    expire_on_commit=False
                )
                
                self.async_session_factory = async_scoped_session(
                    async_session_factory,
                    scopefunc=asyncio.current_task
                )
                
                # 验证新连接
                async with self.async_session_factory() as session:
                    await session.execute(select(1))
                
                logger.info("数据库连接已重新建立")
                return True
            except Exception as reinit_error:
                logger.error(f"重新初始化数据库连接失败: {reinit_error}")
                return False
        except Exception as e:
            logger.error(f"检查数据库连接时发生未知错误: {e}")
            return False

    # ==================== 消息相关方法 ====================
    async def save_message(self,
                           msg:Message
                           ) -> bool:
        """异步保存消息到数据库"""
        # 确保数据库连接有效
        if not await self.ensure_connection():
            logger.error("无法保存消息，数据库连接失败")
            return False
        # 确保content是字符串类型
        async with self.async_session_factory() as session:
            try:
                session.add(msg)
                await session.commit()
                return True
            except SQLAlchemyError as e:
                logger.error(f"保存消息SQL错误: {e}")
                await session.rollback()
                return False
            except Exception as e:
                logger.error(f"保存消息未知错误: {e}")
                await session.rollback()
                return False

    async def get_messages(self,
                           start_time: Optional[datetime] = None,
                           end_time: Optional[datetime] = None,
                           sender_wxid: Optional[str] = None,
                           from_wxid: Optional[str] = None,
                           msg_type: Optional[int] = None,
                           is_group: Optional[bool] = None,
                           limit: int = 100) -> List[Message]:
        """异步查询消息记录"""
        async with self.async_session_factory() as session:
            try:
                query = select(Message).order_by(Message.timestamp.desc()).limit(limit)

                if start_time:
                    query = query.where(Message.timestamp >= start_time)
                if end_time:
                    query = query.where(Message.timestamp <= end_time)
                if sender_wxid:
                    query = query.where(Message.sender_wxid == sender_wxid)
                if from_wxid:
                    query = query.where(Message.from_wxid == from_wxid)
                if msg_type is not None:
                    query = query.where(Message.msg_type == msg_type)
                if is_group is not None:
                    query = query.where(Message.is_group == is_group)

                result = await session.execute(query)
                return result.scalars().all()
            except SQLAlchemyError as e:
                logger.error(f"查询消息SQL错误: {e}")
                return []
            except Exception as e:
                logger.error(f"查询消息未知错误: {e}")
                return []

    async def get_messages_by_chat_id(self, chat_id: str, limit: int = 100) -> List[Message]:
        """根据聊天ID获取消息记录"""
        if not chat_id:
            return []
            
        chat_id = str(chat_id)  # 确保chat_id为字符串类型
        
        async with self.async_session_factory() as session:
            try:
                query = select(Message).where(Message.from_wxid == chat_id)\
                    .order_by(Message.timestamp.desc())\
                    .limit(limit)
                result = await session.execute(query)
                messages = result.scalars().all()
                return list(messages)
            except SQLAlchemyError as e:
                logger.error(f"获取聊天记录SQL错误: {e}")
                return []
            except Exception as e:
                logger.error(f"获取聊天记录未知错误: {e}")
                return []

    # ==================== 群组相关方法 ====================
    async def save_group_info(self, qid: str, group_name: str, owner_wxid: str, 
                             member_count: int, announcement: str = "", member_list: Optional[List[Dict[str, Any]]] = None) -> bool:
        """保存群组信息"""
        # 确保字符串类型
        qid = str(qid) if qid else ""
        group_name = str(group_name) if group_name else ""
        owner_wxid = str(owner_wxid) if owner_wxid else ""
        announcement = str(announcement) if announcement else ""
        
        # 确保member_list不为None
        if member_list is None:
            member_list = []
        
        async with self.async_session_factory() as session:
            try:
                # 检查是否已存在
                existing = await session.execute(select(GroupInfo).where(GroupInfo.qid == qid))
                group_info = existing.scalar_one_or_none()
                
                if group_info:
                    # 更新现有记录
                    group_info.group_name = group_name
                    group_info.owner_wxid = owner_wxid
                    group_info.member_count = member_count
                    group_info.announcement = announcement
                    group_info.member_list = member_list
                    group_info.last_update = datetime.now()
                else:
                    # 创建新记录
                    group_info = GroupInfo(
                        qid=qid,
                        group_name=group_name,
                        owner_wxid=owner_wxid,
                        member_count=member_count,
                        announcement=announcement,
                        member_list=member_list,
                        last_update=datetime.now()
                    )
                    session.add(group_info)
                
                await session.commit()
                return True
            except SQLAlchemyError as e:
                logger.error(f"保存群组信息SQL错误: {e}")
                await session.rollback()
                return False
            except Exception as e:
                logger.error(f"保存群组信息未知错误: {e}")
                await session.rollback()
                return False

    async def get_convertion_id(self,group_id,frend_wxid,is_group,attr_name:str):
        if is_group:
            async with self.async_session_factory() as session:
                result = await session.execute(select(GroupInfo).where(GroupInfo.qid==group_id))
                group_info=result.scalars().one_or_none()
                if group_info and group_info.conversation:
                    return group_info.conversation.get(attr_name)
        else:
            async with self.async_session_factory() as session:
                result = await session.execute(select(FriendInfo).where(FriendInfo.wxid==frend_wxid))
                friend_info=result.scalars().one_or_none()
                if friend_info and friend_info.conversation:
                    return friend_info.conversation.get(attr_name)

    async def update_convertion_id(self,group_id, frend_wxid,is_group,convertion_id:str,attr_name:str):
        if is_group:
            async with self.async_session_factory() as session:
                result = await session.execute(select(GroupInfo).where(GroupInfo.qid==group_id))
                group_info=result.scalars().one_or_none()
                if group_info:
                    if not group_info.conversation:
                        group_info.conversation={}
                    group_info.conversation[attr_name]=convertion_id
                    await session.commit()
        else:
            async with self.async_session_factory() as session:
                result = await session.execute(select(FriendInfo).where(FriendInfo.wxid==frend_wxid))
                friend_info=result.scalars().one_or_none()
                if friend_info:
                    if not friend_info.conversation:
                        friend_info.conversation={}
                    friend_info.conversation[attr_name]=convertion_id
                    await session.commit()


    async def get_group_info(self, group_id: str) -> Optional[dict]:
        """获取群组信息"""
        async with self.async_session_factory() as session:
            try:
                # 这里应该查询群组表，暂时返回基本信息
                return {
                    'group_id': group_id,
                    'group_name': group_id,
                    'member_count': 0,
                    'owner_wxid': '',
                    'announcement': '',
                    'members': []
                }
            except Exception as e:
                logger.error(f"获取群组信息失败: {str(e)}")
                return None

    async def get_all_groups(self) -> List[GroupInfo]:
        """获取所有群组信息"""
        async with self.async_session_factory() as session:
            try:
                result = await session.execute(select(GroupInfo).order_by(GroupInfo.last_update.desc()))
                return result.scalars().all()
            except Exception as e:
                logger.error(f"获取所有群组失败: {str(e)}")
                return []

    # ==================== 好友相关方法 ====================
    async def save_friend_info(self, wxid: str, nickname: str, remark: str = "", 
                              avatar_url: str = "", phone: str = "", email: str = "") -> bool:
        """保存好友信息"""
        async with self.async_session_factory() as session:
            try:
                # 检查是否已存在
                existing = await session.execute(select(FriendInfo).where(FriendInfo.wxid == wxid))
                friend_info = existing.scalar_one_or_none()
                
                if friend_info:
                    # 更新现有记录
                    friend_info.nickname = nickname
                    friend_info.remark = remark
                    friend_info.avatar_url = avatar_url
                    friend_info.phone = phone
                    friend_info.email = email
                    friend_info.last_update = datetime.now()
                else:
                    # 创建新记录
                    friend_info = FriendInfo(
                        wxid=wxid,
                        nickname=nickname,
                        remark=remark,
                        avatar_url=avatar_url,
                        phone=phone,
                        email=email
                    )
                    session.add(friend_info)
                
                await session.commit()
                return True
            except Exception as e:
                logger.error(f"保存好友信息失败: {str(e)}")
                await session.rollback()
                return False

    async def get_friend_info(self, wxid: str) -> Optional[dict]:
        """获取好友信息"""
        async with self.async_session_factory() as session:
            try:
                # 这里应该查询好友表，暂时返回基本信息
                return {
                    'wxid': wxid,
                    'nickname': wxid,
                    'remark': wxid,
                    'avatar_url': ''
                }
            except Exception as e:
                logger.error(f"获取好友信息失败: {str(e)}")
                return None

    async def get_all_friends(self) -> List[Dict[str, Any]]:
        """获取所有好友信息"""
        # 确保数据库连接有效
        if not await self.ensure_connection():
            logger.error("无法获取好友列表，数据库连接失败")
            return []
        
        async with self.async_session_factory() as session:
            try:
                # 查询所有好友
                result = await session.execute(select(FriendInfo))
                friends = result.scalars().all()
                
                # 将ORM对象转换为字典
                friend_list = []
                for friend in friends:
                    friend_list.append({
                        "wxid": friend.wxid,
                        "nickname": friend.nickname,
                        "remark": friend.remark,
                        "avatar_url": friend.avatar_url,
                        "last_update": friend.last_update.isoformat() if friend.last_update else None
                    })
                return friend_list
            except SQLAlchemyError as e:
                logger.error(f"获取所有好友SQL错误: {str(e)}")
                return []
            except Exception as e:
                logger.error(f"获取所有好友失败: {str(e)}")
                return []

    # ==================== 聊天记录相关方法 ====================
    async def save_chat_record(self, chat_id: str, chat_type: str, msg_id: int, 
                              sender_wxid: str, content: str, msg_type: int) -> bool:
        """保存聊天记录"""
        async with self.async_session_factory() as session:
            try:
                chat_record = ChatRecord(
                    chat_id=chat_id,
                    chat_type=chat_type,
                    msg_id=msg_id,
                    sender_wxid=sender_wxid,
                    content=content,
                    msg_type=msg_type
                )
                session.add(chat_record)
                await session.commit()
                return True
            except Exception as e:
                logger.error(f"保存聊天记录失败: {str(e)}")
                await session.rollback()
                return False

    async def get_chat_records(self, chat_id: str, limit: int = 100) -> List[ChatRecord]:
        """获取聊天记录"""
        async with self.async_session_factory() as session:
            try:
                result = await session.execute(
                    select(ChatRecord)
                    .where(ChatRecord.chat_id == chat_id)
                    .order_by(ChatRecord.timestamp.desc())
                    .limit(limit)
                )
                return result.scalars().all()
            except Exception as e:
                logger.error(f"获取聊天记录失败: {str(e)}")
                return []

    # ==================== 聊天记录列表相关方法 ====================
    async def get_chat_list(self, chat_type: str = None) -> List[dict]:
        """获取聊天记录列表，按聊天对象分组统计"""
        async with self.async_session_factory() as session:
            try:
                # 构建查询
                if chat_type == 'group':
                    # 查询群聊记录
                    query = select(Message)\
                        .where(Message.is_group == True)\
                        .order_by(Message.timestamp.desc())
                elif chat_type == 'private':
                    # 查询私聊记录
                    query = select(Message)\
                        .where(Message.is_group == False)\
                        .order_by(Message.timestamp.desc())
                else:
                    # 查询所有记录
                    query = select(Message)\
                        .order_by(Message.timestamp.desc())
                
                result = await session.execute(query)
                messages = result.scalars().all()
                
                # 按聊天对象分组统计
                chat_stats = {}
                for msg in messages:
                    chat_id = msg.from_wxid
                    if not chat_id:  # 跳过空的聊天ID
                        continue
                        
                    if chat_id not in chat_stats:
                        chat_stats[chat_id] = {
                            'chat_id': chat_id,
                            'chat_name': getattr(msg,'from_nickname',''),
                            'chat_type': 'group' if getattr(msg, 'is_group', False) else 'private',
                            'last_message': getattr(msg,'org_content',''),
                            'at_names': getattr(msg,'at_names',''),
                            'sender_nickname':getattr(msg,'sender_nickname',''),
                            'last_time': msg.timestamp,
                            'message_count': 0,
                            'unread_count': 0
                        }
                    chat_stats[chat_id]['message_count'] += 1
                
                # 转换为列表并按最后消息时间排序
                chat_list = list(chat_stats.values())
                chat_list.sort(key=lambda x: x['last_time'], reverse=True)
                
                return chat_list
            except Exception as e:
                logger.error(f"获取聊天列表失败: {str(e)}")
                return []

    async def get_chat_detail(self, chat_id: str, limit: int = 100) -> List[dict]:
        """获取指定聊天的详细消息记录"""
        async with self.async_session_factory() as session:
            try:
                query = select(Message)\
                    .where(Message.from_wxid == chat_id)\
                    .order_by(Message.timestamp.desc())\
                    .limit(limit)
                
                result = await session.execute(query)
                messages = result.scalars().all()
                
                # 转换为字典格式
                message_list = []
                for msg in messages:
                    message_list.append({
                        'id': msg.id,
                        'msg_id': msg.msg_id,
                        'sender_wxid': msg.sender_wxid,
                        'content': msg.content,
                        'msg_type': msg.msg_type,
                        'timestamp': msg.timestamp.isoformat(),
                        'is_group': msg.is_group
                    })
                
                return message_list
            except Exception as e:
                logger.error(f"获取聊天详情失败: {str(e)}")
                return []

    async def get_message_statistics(self) -> dict:
        """获取消息统计信息"""
        async with self.async_session_factory() as session:
            try:
                # 总消息数
                total_query = select(Message.id)
                total_result = await session.execute(total_query)
                total_messages = len(total_result.all())
                
                # 群消息数
                group_query = select(Message.id).where(Message.is_group == True)
                group_result = await session.execute(group_query)
                group_messages = len(group_result.all())
                
                # 私聊消息数
                private_query = select(Message.id).where(Message.is_group == False)
                private_result = await session.execute(private_query)
                private_messages = len(private_result.all())
                
                # 今日消息数
                today = datetime.now().replace(hour=0, minute=0, second=0, microsecond=0)
                today_query = select(Message.id).where(Message.timestamp >= today)
                today_result = await session.execute(today_query)
                today_messages = len(today_result.all())
                
                return {
                    'total_messages': total_messages,
                    'group_messages': group_messages,
                    'private_messages': private_messages,
                    'today_messages': today_messages
                }
            except Exception as e:
                logger.error(f"获取消息统计失败: {str(e)}")
                return {}

    # ==================== 群成员相关方法 ====================
    async def save_group_member(self, group_id: str, member_wxid: str, member_nickname: str = "", 
                               member_avatar: str = "", join_time: datetime = None) -> bool:
        """保存群成员信息"""
        async with self.async_session_factory() as session:
            try:
                # 这里可以扩展为独立的群成员表
                # 暂时更新到群组信息的member_list字段
                group_info = await self.get_group_info(group_id)
                if group_info:
                    member_list = group_info.member_list or []
                    
                    # 查找是否已存在该成员
                    member_exists = False
                    for member in member_list:
                        if member.get('wxid') == member_wxid:
                            member['nickname'] = member_nickname
                            member['avatar'] = member_avatar
                            member['join_time'] = join_time.isoformat() if join_time else None
                            member_exists = True
                            break
                    
                    if not member_exists:
                        member_list.append({
                            'wxid': member_wxid,
                            'nickname': member_nickname,
                            'avatar': member_avatar,
                            'join_time': join_time.isoformat() if join_time else None
                        })
                    
                    # 更新群组信息
                    group_info.member_list = member_list
                    group_info.last_update = datetime.now()
                    await session.commit()
                    return True
                
                return False
            except Exception as e:
                logger.error(f"保存群成员信息失败: {str(e)}")
                await session.rollback()
                return False

    async def get_group_members(self, group_id: str) -> List[dict]:
        """获取群成员列表"""
        group_info = await self.get_group_info(group_id)
        if group_info and group_info.member_list:
            return group_info.member_list
        return []

    # ==================== 好友信息扩展方法 ====================
    async def update_friend_avatar(self, wxid: str, avatar_url: str) -> bool:
        """更新好友头像"""
        async with self.async_session_factory() as session:
            try:
                friend = await session.execute(select(FriendInfo).where(FriendInfo.wxid == wxid))
                friend_info = friend.scalar_one_or_none()
                
                if friend_info:
                    friend_info.avatar_url = avatar_url
                    friend_info.last_update = datetime.now()
                    await session.commit()
                    return True
                return False
            except Exception as e:
                logger.error(f"更新好友头像失败: {str(e)}")
                await session.rollback()
                return False

    async def search_friends(self, keyword: str) -> List[FriendInfo]:
        """搜索好友"""
        async with self.async_session_factory() as session:
            try:
                query = select(FriendInfo).where(
                    (FriendInfo.wxid.contains(keyword)) |
                    (FriendInfo.nickname.contains(keyword)) |
                    (FriendInfo.remark.contains(keyword))
                )
                result = await session.execute(query)
                return result.scalars().all()
            except Exception as e:
                logger.error(f"搜索好友失败: {str(e)}")
                return []

    async def close(self):
        """关闭数据库连接"""
        if self.cleanup_task:
            self.cleanup_task.cancel()
            try:
                await self.cleanup_task
            except asyncio.CancelledError:
                pass
        
        # 关闭连接池
        if hasattr(self, 'engine'):
            await self.engine.dispose()
            logger.info("数据库连接已关闭")

    async def cleanup_messages(self):
        """清理过期消息（保留最近30天）"""
        async with self.async_session_factory() as session:
            try:
                # 计算30天前的日期
                cutoff_date = datetime.now() - timedelta(days=30)
                
                # 删除30天前的消息
                stmt = delete(Message).where(Message.timestamp < cutoff_date)
                result = await session.execute(stmt)
                deleted_count = result.rowcount
                
                await session.commit()
                logger.info(f"已清理 {deleted_count} 条过期消息")
                return deleted_count
            except SQLAlchemyError as e:
                logger.error(f"清理消息SQL错误: {e}")
                await session.rollback()
                return 0
            except Exception as e:
                logger.error(f"清理消息未知错误: {e}")
                await session.rollback()
                return 0

    # 上下文管理器支持
    async def __aenter__(self):
        # 初始化数据库
        if not self.initialized:
            await self.initialize()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self.close()
